{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d96a9939",
   "metadata": {},
   "source": [
    "# 다중 벡터저장소 검색기(MultiVectorRetriever)\n",
    "\n",
    "LangChain에서는 문서를 다양한 상황에서 효율적으로 쿼리할 수 있는 특별한 기능, 바로 `MultiVectorRetriever`를 제공합니다. 이 기능을 사용하면 문서를 여러 벡터로 저장하고 관리할 수 있어, 정보 검색의 정확도와 효율성을 대폭 향상시킬 수 있습니다. \n",
    "\n",
    "`MultiVectorRetriever`를 활용해 문서당 여러 벡터를 생성하는 몇 가지 방법을 살펴보겠습니다.\n",
    "\n",
    "**문서당 여러 벡터 생성 방법 소개**\n",
    "\n",
    "1. **작은 청크 생성**: 문서를 더 작은 단위로 나눈 후, 각 청크에 대해 별도의 임베딩을 생성합니다. 이 방식을 사용하면 문서의 특정 부분에 좀 더 세심한 주의를 기울일 수 있습니다. 이 과정은 `ParentDocumentRetriever`를 통해 구현할 수 있어, 세부 정보에 대한 탐색이 용이해집니다.\n",
    "\n",
    "2. **요약 임베딩**: 각 문서의 요약을 생성하고, 이 요약으로부터 임베딩을 만듭니다. 이 요약 임베딩은 문서의 핵심 내용을 신속하게 파악하는 데 큰 도움이 됩니다. 문서 전체를 분석하는 대신 핵심적인 요약 부분만을 활용하여 효율성을 극대화할 수 있습니다.\n",
    "\n",
    "3. **가설 질문 활용**: 각 문서에 대해 적합한 가설 질문을 만들고, 이 질문에 기반한 임베딩을 생성합니다. 특정 주제나 내용에 대해 깊이 있는 탐색을 원할 때 이 방법이 유용합니다. 가설 질문은 문서의 내용을 다양한 관점에서 접근하게 해주며, 더 광범위한 이해를 가능하게 합니다.\n",
    "\n",
    "4. **수동 추가 방식**: 사용자가 문서 검색 시 고려해야 할 특정 질문이나 쿼리를 직접 추가할 수 있습니다. 이 방법을 통해 사용자는 검색 과정에서 보다 세밀한 제어를 할 수 있으며, 자신의 요구 사항에 맞춘 맞춤형 검색이 가능해집니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3645c156",
   "metadata": {},
   "source": [
    "## 실습에 활용한 문서\n",
    "\n",
    "소프트웨어정책연구소(SPRi) - 2023년 12월호\n",
    "\n",
    "- 저자: 유재흥(AI정책연구실 책임연구원), 이지수(AI정책연구실 위촉연구원)\n",
    "- 링크: https://spri.kr/posts/view/23669\n",
    "- 파일명: `SPRI_AI_Brief_2023년12월호_F.pdf`\n",
    "\n",
    "**참고**: 위의 파일은 `data` 폴더 내에 다운로드 받으세요"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18a2a1b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26eac101",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH10-Retriever\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b574ec5",
   "metadata": {},
   "source": [
    "텍스트 파일에서 데이터를 로드하고, 로드된 문서들을 지정된 크기로 분할하는 전처리 과정을 수행합니다.\n",
    "\n",
    "분할된 문서들은 추후 벡터화 및 검색 등의 작업에 사용될 수 있습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d6c52a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PyMuPDFLoader\n",
    "\n",
    "loader = PyMuPDFLoader(\"data/SPRI_AI_Brief_2023년12월호_F.pdf\")\n",
    "docs = loader.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "453afd1e",
   "metadata": {},
   "source": [
    "데이터로부터 로드한 원본 도큐먼트는 `docs` 변수에 담았습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a50910",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(docs[5].page_content[:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b671d6e",
   "metadata": {},
   "source": [
    "## Chunk + 원본 문서 검색\n",
    "\n",
    "대용량 정보를 검색하는 경우, 더 작은 단위로 정보를 임베딩하는 것이 유용할 수 있습니다.\n",
    "\n",
    "`MultiVectorRetriever` 를 통해 문서를 여러 벡터로 저장하고 관리할 수 있습니다.\n",
    "\n",
    "`docstore` 에 원본 문서를 저장하고, `vectorstore` 에 임베딩된 문서를 저장합니다. \n",
    "\n",
    "이로써 문서를 더 작은 단위로 나누어 더 정확한 검색이 가능해집니다. 때에 따라서는 원본 문서의 내용을 조회할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc51c01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 자식 청크를 인덱싱하는 데 사용할 벡터 저장소\n",
    "import uuid\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"small_bigger_chunks\",\n",
    "    embedding_function=OpenAIEmbeddings(model=\"text-embedding-3-small\"),\n",
    ")\n",
    "# 부모 문서의 저장소 계층\n",
    "store = InMemoryStore()\n",
    "\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "# 검색기 (시작 시 비어 있음)\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "\n",
    "# 문서 ID를 생성합니다.\n",
    "doc_ids = [str(uuid.uuid4()) for _ in docs]\n",
    "\n",
    "# 두개의 생성된 id를 확인합니다.\n",
    "doc_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aa3e771",
   "metadata": {},
   "source": [
    "여기서 큰 청크로 분할하기 위한 `parent_text_splitter`\n",
    "\n",
    "더 작은 청크로 분할하기 위한 `child_text_splitter` 를 정의합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc95363",
   "metadata": {},
   "outputs": [],
   "source": [
    "# RecursiveCharacterTextSplitter 객체를 생성합니다.\n",
    "parent_text_splitter = RecursiveCharacterTextSplitter(chunk_size=600)\n",
    "\n",
    "# 더 작은 청크를 생성하는 데 사용할 분할기\n",
    "child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8e48a98",
   "metadata": {},
   "source": [
    "더 큰 Chunk인 Parent 문서를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4437c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "parent_docs = []\n",
    "\n",
    "for i, doc in enumerate(docs):\n",
    "    # 현재 문서의 ID를 가져옵니다.\n",
    "    _id = doc_ids[i]\n",
    "    # 현재 문서를 하위 문서로 분할\n",
    "    parent_doc = parent_text_splitter.split_documents([doc])\n",
    "\n",
    "    for _doc in parent_doc:\n",
    "        # metadata에 문서 ID 를 저장\n",
    "        _doc.metadata[id_key] = _id\n",
    "    parent_docs.extend(parent_doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c5b0e5",
   "metadata": {},
   "source": [
    "`parent_docs` 에 기입된 `doc_id` 를 확인합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d7ff8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 생성된 Parent 문서의 메타데이터를 확인합니다.\n",
    "parent_docs[0].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f680da0",
   "metadata": {},
   "source": [
    "상대적으로 더 작은 Chunk인 Child 문서를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e56afe1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "child_docs = []\n",
    "for i, doc in enumerate(docs):\n",
    "    # 현재 문서의 ID를 가져옵니다.\n",
    "    _id = doc_ids[i]\n",
    "    # 현재 문서를 하위 문서로 분할\n",
    "    child_doc = child_text_splitter.split_documents([doc])\n",
    "    for _doc in child_doc:\n",
    "        # metadata에 문서 ID 를 저장\n",
    "        _doc.metadata[id_key] = _id\n",
    "    child_docs.extend(child_doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bafe5cb5",
   "metadata": {},
   "source": [
    "`child_docs` 에 기입된 `doc_id` 를 확인합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80857992",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 생성된 Child 문서의 메타데이터를 확인합니다.\n",
    "child_docs[0].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a346a79",
   "metadata": {},
   "source": [
    "각각 분할된 청크의 수를 확인합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dfd9565",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"분할된 parent_docs의 개수: {len(parent_docs)}\")\n",
    "print(f\"분할된 child_docs의 개수: {len(child_docs)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6da6ab64",
   "metadata": {},
   "source": [
    "벡터저장소에 새롭게 생성한 작게 쪼개진 하위문서 집합을 추가합니다.\n",
    "\n",
    "다음으로는 상위 문서는 생성한 UUID 와 맵핑하여 `docstore` 에 추가합니다.\n",
    "\n",
    "- `mset()` 메서드를 통해 문서 ID와 문서 내용을 key-value 쌍으로 문서 저장소에 저장합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a291a760",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 벡터 저장소에 parent + child 문서를 추가\n",
    "retriever.vectorstore.add_documents(parent_docs)\n",
    "retriever.vectorstore.add_documents(child_docs)\n",
    "\n",
    "# docstore 에 원본 문서를 저장\n",
    "retriever.docstore.mset(list(zip(doc_ids, docs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92a644c1",
   "metadata": {},
   "source": [
    "유사도 검색을 수행합니다. 가장 유사도가 높은 첫 번째 문서 조각을 출력합니다.\n",
    "\n",
    "여기서 `retriever.vectorstore.similarity_search` 메서드는 child + parent 문서 chunk 내에서 검색을 수행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d9ba0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# vectorstore의 유사도 검색을 수행합니다.\n",
    "relevant_chunks = retriever.vectorstore.similarity_search(\n",
    "    \"삼성전자가 만든 생성형 AI 의 이름은?\"\n",
    ")\n",
    "print(f\"검색된 문서의 개수: {len(relevant_chunks)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5eb7fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for chunk in relevant_chunks:\n",
    "    print(chunk.page_content, end=\"\\n\\n\")\n",
    "    print(\">\" * 100, end=\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "598f6f1a",
   "metadata": {},
   "source": [
    "이번에는 `retriever.invoke()` 메서드를 사용하여 쿼리를 실행합니다.\n",
    "\n",
    "`retriever.invoke()` 메서드는 원본 문서의 전체 내용을 검색합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a29f9b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "relevant_docs = retriever.invoke(\"삼성전자가 만든 생성형 AI 의 이름은?\")\n",
    "print(f\"검색된 문서의 개수: {len(relevant_docs)}\", end=\"\\n\\n\")\n",
    "print(\"=\" * 100, end=\"\\n\\n\")\n",
    "print(relevant_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acc30896",
   "metadata": {},
   "source": [
    "리트리버(retriever)가 벡터 데이터베이스에서 기본적으로 수행하는 검색 유형은 유사도 검색입니다.\n",
    "\n",
    "LangChain Vector Stores는 [Max Marginal Relevance](https://api.python.langchain.com/en/latest/vectorstores/langchain_core.vectorstores.VectorStore.html#langchain_core.vectorstores.VectorStore.max_marginal_relevance_search)를 통한 검색도 지원하므로, 이를 대신 사용하고 싶다면 다음과 같이 `search_type` 속성을 설정하면 됩니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac718689",
   "metadata": {},
   "source": [
    "- `retriever` 객체의 `search_type` 속성을 `SearchType.mmr`로 설정합니다.\n",
    "  - 이는 검색 시 MMR(Maximal Marginal Relevance) 알고리즘을 사용하도록 지정하는 것입니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c545fc19",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.multi_vector import SearchType\n",
    "\n",
    "# 검색 유형을 MMR(Maximal Marginal Relevance)로 설정\n",
    "retriever.search_type = SearchType.mmr\n",
    "\n",
    "# 관련 문서 전체를 검색\n",
    "print(retriever.invoke(\"삼성전자가 만든 생성형 AI 의 이름은?\")[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07d08019",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.multi_vector import SearchType\n",
    "\n",
    "# 검색 유형을 similarity_score_threshold로 설정\n",
    "retriever.search_type = SearchType.similarity_score_threshold\n",
    "retriever.search_kwargs = {\"score_threshold\": 0.3}\n",
    "\n",
    "# 관련 문서 전체를 검색\n",
    "print(retriever.invoke(\"삼성전자가 만든 생성형 AI 의 이름은?\")[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d203555f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.multi_vector import SearchType\n",
    "\n",
    "# 검색 유형을 similarity로 설정, k값을 1로 설정\n",
    "retriever.search_type = SearchType.similarity\n",
    "retriever.search_kwargs = {\"k\": 1}\n",
    "\n",
    "# 관련 문서 전체를 검색\n",
    "print(len(retriever.invoke(\"삼성전자가 만든 생성형 AI 의 이름은?\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d961b6c",
   "metadata": {},
   "source": [
    "## 요약본(summary)을 벡터저장소에 저장\n",
    "\n",
    "요약은 종종 청크(chunk)의 내용을 보다 정확하게 추출할 수 있어 더 나은 검색 결과를 얻을 수 있습니다.\n",
    "\n",
    "여기서는 요약을 생성하는 방법과 이를 임베딩하는 방법에 대해 설명합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7c44d7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PDF 파일을 로드하고 텍스트를 분할하기 위한 라이브러리 임포트\n",
    "from langchain_community.document_loaders import PyMuPDFLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "# PDF 파일 로더 초기화\n",
    "loader = PyMuPDFLoader(\"data/SPRI_AI_Brief_2023년12월호_F.pdf\")\n",
    "\n",
    "# 텍스트 분할\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=50)\n",
    "\n",
    "# PDF 파일 로드 및 텍스트 분할 실행\n",
    "split_docs = loader.load_and_split(text_splitter)\n",
    "\n",
    "# 분할된 문서의 개수 출력\n",
    "print(f\"분할된 문서의 개수: {len(split_docs)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a161ea24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.documents import Document\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "summary_chain = (\n",
    "    {\"doc\": lambda x: x.page_content}\n",
    "    # 문서 요약을 위한 프롬프트 템플릿 생성\n",
    "    | ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", \"You are an expert in summarizing documents in Korean.\"),\n",
    "            (\n",
    "                \"user\",\n",
    "                \"Summarize the following documents in 3 sentences in bullet points format.\\n\\n{doc}\",\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    # OpenAI의 ChatGPT 모델을 사용하여 요약 생성\n",
    "    | ChatOpenAI(temperature=0, model=\"gpt-4o-mini\")\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16b79165",
   "metadata": {},
   "source": [
    "`chain.batch` 메서드를 사용하여 `docs` 리스트의 문서들을 일괄 요약합니다.\n",
    "- 여기서 `max_concurrency` 매개변수를 10 으로 설정하여 최대 10개의 문서를 동시에 처리할 수 있도록 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7fe3a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서 배치 처리\n",
    "summaries = summary_chain.batch(split_docs, {\"max_concurrency\": 10})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45363e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(summaries)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1a11208",
   "metadata": {},
   "source": [
    "요약된 내용을 출력하여 결과를 확인합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "436bd00a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 원본 문서의 내용을 출력합니다.\n",
    "print(split_docs[33].page_content, end=\"\\n\\n\")\n",
    "# 요약을 출력합니다.\n",
    "print(\"[요약]\")\n",
    "print(summaries[33])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04da5865",
   "metadata": {},
   "source": [
    "`Chroma` 벡터 저장소를 초기화하여 자식 청크(child chunks)를 인덱싱합니다. 이때 `OpenAIEmbeddings`를 임베딩 함수로 사용합니다.\n",
    "\n",
    "- 문서 ID를 나타내는 키로 `\"doc_id\"`를 사용합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe2f51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "# 요약 정보를 저장할 벡터 저장소를 생성합니다.\n",
    "summary_vectorstore = Chroma(\n",
    "    collection_name=\"summaries\",\n",
    "    embedding_function=OpenAIEmbeddings(model=\"text-embedding-3-small\"),\n",
    ")\n",
    "\n",
    "# 부모 문서를 저장할 저장소를 생성합니다.\n",
    "store = InMemoryStore()\n",
    "\n",
    "# 문서 ID를 저장할 키 이름을 지정합니다.\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "# 검색기를 초기화합니다. (시작 시 비어 있음)\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=summary_vectorstore,  # 벡터 저장소\n",
    "    byte_store=store,  # 바이트 저장소\n",
    "    id_key=id_key,  # 문서 ID 키\n",
    ")\n",
    "# 문서 ID를 생성합니다.\n",
    "doc_ids = [str(uuid.uuid4()) for _ in split_docs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9d0e2c6",
   "metadata": {},
   "source": [
    "요약된 문서와 메타데이터(여기서는 생성한 요약본에 대한 `Document ID` 입니다)를 저장합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cec99148",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_docs = [\n",
    "    # 요약된 내용을 페이지 콘텐츠로 하고, 문서 ID를 메타데이터로 포함하는 Document 객체를 생성합니다.\n",
    "    Document(page_content=s, metadata={id_key: doc_ids[i]})\n",
    "    for i, s in enumerate(summaries)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb07ac1",
   "metadata": {},
   "source": [
    "요약본의 문서의 개수는 원본 문서의 개수와 일치합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d36e4d50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 요약본의 문서의 개수\n",
    "len(summary_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4400de8c",
   "metadata": {},
   "source": [
    "- `retriever.vectorstore.add_documents(summary_docs)`를 통해 `summary_docs`를 벡터 저장소에 추가합니다.\n",
    "- `retriever.docstore.mset(list(zip(doc_ids, docs)))`를 사용하여 `doc_ids`와 `docs`를 매핑하여 문서 저장소에 저장합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7ce3a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever.vectorstore.add_documents(\n",
    "    summary_docs\n",
    ")  # 요약된 문서를 벡터 저장소에 추가합니다.\n",
    "\n",
    "# 문서 ID와 문서를 매핑하여 문서 저장소에 저장합니다.\n",
    "retriever.docstore.mset(list(zip(doc_ids, split_docs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9158831",
   "metadata": {},
   "source": [
    "`vectorstore` 객체의 `similarity_search` 메서드를 사용하여 유사도 검색을 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e3d643",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 유사도 검색을 수행합니다.\n",
    "result_docs = summary_vectorstore.similarity_search(\n",
    "    \"삼성전자가 만든 생성형 AI 의 이름은?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7cb5a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1개의 결과 문서를 출력합니다.\n",
    "print(result_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bd64a80",
   "metadata": {},
   "source": [
    "`retriever` 객체의 `invoke()` 사용하여 질문과 관련된 문서를 검색합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5a9c432",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 관련된 문서를 검색하여 가져옵니다.\n",
    "retrieved_docs = retriever.invoke(\"삼성전자가 만든 생성형 AI 의 이름은?\")\n",
    "print(retrieved_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e3d5f1b",
   "metadata": {},
   "source": [
    "## 가설 쿼리(Hypothetical Queries) 를 활용하여 문서 내용 탐색\n",
    "\n",
    "LLM은 특정 문서에 대해 가정할 수 있는 질문 목록을 생성하는 데에도 사용될 수 있습니다.\n",
    "\n",
    "이렇게 생성된 질문들은 임베딩(embedding)될 수 있으며, 이를 통해 문서의 내용을 더욱 깊이 있게 탐색하고 이해할 수 있습니다.\n",
    "\n",
    "가정 질문 생성은 문서의 주요 주제와 개념을 파악하는 데 도움이 되며, 독자들이 문서 내용에 대해 더 많은 궁금증을 갖도록 유도할 수 있습니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7933e03",
   "metadata": {},
   "source": [
    "아래는 `Function Calling` 을 활용하여 가설 질문을 생성하는 예제입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5baf8ff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "functions = [\n",
    "    {\n",
    "        \"name\": \"hypothetical_questions\",  # 함수의 이름을 지정합니다.\n",
    "        \"description\": \"Generate hypothetical questions\",  # 함수에 대한 설명을 작성합니다.\n",
    "        \"parameters\": {  # 함수의 매개변수를 정의합니다.\n",
    "            \"type\": \"object\",  # 매개변수의 타입을 객체로 지정합니다.\n",
    "            \"properties\": {  # 객체의 속성을 정의합니다.\n",
    "                \"questions\": {  # 'questions' 속성을 정의합니다.\n",
    "                    \"type\": \"array\",  # 'questions'의 타입을 배열로 지정합니다.\n",
    "                    \"items\": {\n",
    "                        \"type\": \"string\"\n",
    "                    },  # 배열의 요소 타입을 문자열로 지정합니다.\n",
    "                },\n",
    "            },\n",
    "            \"required\": [\"questions\"],  # 필수 매개변수로 'questions'를 지정합니다.\n",
    "        },\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64aa557a",
   "metadata": {},
   "source": [
    "`ChatPromptTemplate`을 사용하여 주어진 문서를 기반으로 3개의 가상 질문을 생성하는 프롬프트 템플릿을 정의합니다.\n",
    "\n",
    "- `functions`와 `function_call`을 설정하여 가상 질문 생성 함수를 호출합니다.\n",
    "- `JsonKeyOutputFunctionsParser`를 사용하여 생성된 가상 질문을 파싱하고, `questions` 키에 해당하는 값을 추출합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cb4be57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain.output_parsers.openai_functions import JsonKeyOutputFunctionsParser\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "hypothetical_query_chain = (\n",
    "    {\"doc\": lambda x: x.page_content}\n",
    "    # 아래 문서를 사용하여 답변할 수 있는 가상의 질문을 정확히 3개 생성하도록 요청합니다. 이 숫자는 조정될 수 있습니다.\n",
    "    | ChatPromptTemplate.from_template(\n",
    "        \"Generate a list of exactly 3 hypothetical questions that the below document could be used to answer. \"\n",
    "        \"Potential users are those interested in the AI industry. Create questions that they would be interested in. \"\n",
    "        \"Output should be written in Korean:\\n\\n{doc}\"\n",
    "    )\n",
    "    | ChatOpenAI(max_retries=0, model=\"gpt-4o-mini\").bind(\n",
    "        functions=functions, function_call={\"name\": \"hypothetical_questions\"}\n",
    "    )\n",
    "    # 출력에서 \"questions\" 키에 해당하는 값을 추출합니다.\n",
    "    | JsonKeyOutputFunctionsParser(key_name=\"questions\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b789211",
   "metadata": {},
   "source": [
    "문서에 대한 답변을 출력합니다.\n",
    "\n",
    "- 출력은 생성한 3개의 가설 쿼리(Hypothetical Queries) 가 담겨 있습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9cda3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 주어진 문서에 대해 체인을 실행합니다.\n",
    "hypothetical_query_chain.invoke(split_docs[33])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853ff912",
   "metadata": {},
   "source": [
    "`chain.batch` 메서드를 사용하여 `split_docs` 데이터에 대해 동시에 여러 개의 요청을 처리합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac642190",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서 목록에 대해 가설 질문을 배치 생성\n",
    "hypothetical_questions = hypothetical_query_chain.batch(\n",
    "    split_docs, {\"max_concurrency\": 10}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02840bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "hypothetical_questions[33]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f364ded0",
   "metadata": {},
   "source": [
    "아래는 이전에 진행했던 방식과 동일하게 생성한 가설 쿼리(Hypothetical Queries) 를 벡터저장소에 저장하는 과정입니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c44404c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 자식 청크를 인덱싱하는 데 사용할 벡터 저장소\n",
    "hypothetical_vectorstore = Chroma(\n",
    "    collection_name=\"hypo-questions\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "# 부모 문서의 저장소 계층\n",
    "store = InMemoryStore()\n",
    "\n",
    "id_key = \"doc_id\"\n",
    "# 검색기 (시작 시 비어 있음)\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=hypothetical_vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "doc_ids = [str(uuid.uuid4()) for _ in split_docs]  # 문서 ID 생성"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8f8a6b7",
   "metadata": {},
   "source": [
    "`question_docs` 리스트에 메타데이터(문서 ID) 를 추가합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7cbb65e",
   "metadata": {},
   "outputs": [],
   "source": [
    "question_docs = []\n",
    "# hypothetical_questions 저장\n",
    "for i, question_list in enumerate(hypothetical_questions):\n",
    "    question_docs.extend(\n",
    "        # 질문 리스트의 각 질문에 대해 Document 객체를 생성하고, 메타데이터에 해당 질문의 문서 ID를 포함시킵니다.\n",
    "        [Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "480a8e5f",
   "metadata": {},
   "source": [
    "가설 쿼리를 문서에 추가하고, 원본 문서를 `docstore` 에 추가합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "105ecf34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hypothetical_questions 문서를 벡터 저장소에 추가합니다.\n",
    "retriever.vectorstore.add_documents(question_docs)\n",
    "\n",
    "# 문서 ID와 문서를 매핑하여 문서 저장소에 저장합니다.\n",
    "retriever.docstore.mset(list(zip(doc_ids, split_docs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "badb5835",
   "metadata": {},
   "source": [
    "`vectorstore` 객체의 `similarity_search` 메서드를 사용하여 유사도 검색을 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca9afb0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 유사한 문서를 벡터 저장소에서 검색합니다.\n",
    "result_docs = hypothetical_vectorstore.similarity_search(\n",
    "    \"삼성전자가 만든 생성형 AI 의 이름은?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2ce41e6",
   "metadata": {},
   "source": [
    "아래는 유사도 검색 결과입니다.\n",
    "\n",
    "여기서는 생성한 가설 쿼리만 추가해 놓은 상태이기 때문에, 생성한 가설 쿼리 중 유사도가 가장 높은 문서를 반환합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c33bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 유사도 검색 결과를 출력합니다.\n",
    "for doc in result_docs:\n",
    "    print(doc.page_content)\n",
    "    print(doc.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88722c7",
   "metadata": {},
   "source": [
    "`retriever` 객체의 `invoke` 메서드를 사용하여 쿼리와 관련된 문서를 검색합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d106a8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 관련된 문서를 검색하여 가져옵니다.\n",
    "retrieved_docs = retriever.invoke(result_docs[1].page_content)\n",
    "\n",
    "# 검색된 문서를 출력합니다.\n",
    "for doc in retrieved_docs:\n",
    "    print(doc.page_content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
